#include <lcir/ADT/Range.h>
#include <lcir/IR/Argument.h>
#include <lcir/IR/Constant.h>
#include <lcir/IR/Expression.h>
#include <lcir/IR/Lambda.h>
#include <lcir/IR/LambdaType.h>
#include <lcir/IR/VReg.h>
#include <lcir/IR/Value.h>
#include <lcir/IR/ValuePrinter.h>

#include <cstddef>
#include <sstream>
#include <string>
using namespace std::string_literals;
namespace lcir {
static void printTab(std::stringstream &ss, size_t depth) {
  for (size_t i = 0; i < depth; ++i) ss << " ";
}

static void mapInit(Value *root, std::map<const Value *, size_t> &valToIdx) {
  size_t idx = 0;
  auto *lambda = root->dyn_cast<Lambda>();
  if (!lambda) {
    valToIdx[root] = idx;
    return;
  }
  lambda->walk([&](Expression *expr) {
    for (auto *result : expr->getResults()) {
      valToIdx[result] = idx++;
    }
  });
}
ValuePrinter::ValuePrinter(const Value *value) {
  if (!value) return;
  root = value->getRoot();
  mapInit(root, valueToIdx);
}

std::string ValuePrinter::getVregIdx(const Value *value) {
  auto it = valueToIdx.find(value);
  std::stringstream ss;
  ss << "%";
  if (it == valueToIdx.end()) {
    ss << "unknown";
  } else {
    ss << it->second;
    if (!value->getInfo().empty()) ss << "(" << value->getInfo() << ")";
  }
  return ss.str();
}

std::string ValuePrinter::print(const Constant *constant) { return constant->getConstStr(); }
std::string ValuePrinter::print(const Argument *argument) {
  std::stringstream ss;
  std::cout << ss.str();
  auto it = argToName.find(argument);
  if (it == argToName.end()) {
    ss << "unknown_arg";
  } else
    ss << it->second;
  return ss.str();
}

std::string ValuePrinter::print(const VReg *vreg) {
  std::string name = getVregIdx(vreg);
  return name;
}

std::string ValuePrinter::print(const Lambda *lambda) {
  std::stringstream ss;
  auto depth = lambda->getDepth();

  // 1. 打印参数
  // 如果这个lambda在操作数上，多打印个换行符
  if (lambda->getDefExpr() && !lambda->getExprs().empty()) {
    ss << "\n";
    printTab(ss, ++depth);
  }
  ss << "(";
  auto args = lambda->getArgs();
  for (const auto &[idx, arg] : IndexRange(args)) {
    argToName[arg] = R"(%arg)" + std::to_string(idx) + "("s + arg->getInfo() + ")"s;
    ss << argToName[arg];
    ss << ": " << arg->getType()->toString();
    if (idx != args.size() - 1) ss << ", ";
  }
  ss << ")";
  // 2. 打印返回值
  auto *lambdaTy = lambda->getLambdaType();
  auto retTy = lambdaTy->getResultType();
  if (retTy) {
    ss << ": " << retTy->toString();
  }
  ss << " => {";
  // 递归打印expression
  if (!lambda->getExprs().empty()) ss << "\n";
  for (auto &expr : lambda->getExprs()) {
    ss << print(&expr) << "\n";
  }
  printTab(ss, depth);
  ss << "}";
  return ss.str();
}

std::string ValuePrinter::print(const Value *value) {
  if (auto *constant = value->dyn_cast<Constant>()) {
    return print(constant);
  } else if (auto *argument = value->dyn_cast<Argument>()) {
    return print(argument);
  } else if (auto *vreg = value->dyn_cast<VReg>()) {
    return print(vreg);
  } else if (auto *lambda = value->dyn_cast<Lambda>()) {
    return print(lambda);
  }
  return "unknown_value";
}

std::string ValuePrinter::print(const Expression *expr) {
  std::stringstream ss;
  auto depth = expr->getDepth() + 1;
  // if (auto *lambda = expr->getLambda()) {
  //   if (lambda->getDefExpr()) depth++;
  // }
  printTab(ss, depth);
  auto rets = expr->getResults();
  auto operands = expr->getOperands();
  if (rets.size() > 1) ss << "(";
  for (auto [idx, ret] : IndexRange(rets)) {
    ss << print(ret);
    if (idx != rets.size() - 1) ss << ", ";
  }
  if (rets.size() > 1) ss << " )";
  if (!rets.empty()) ss << " = ";
  ss << expr->getOp()->getName();

  ss << "(";
  // 然后打印operand
  for (auto [idx, opd] : IndexRange(operands)) {
    ss << print(opd->getValue());
    if (idx != operands.size() - 1) ss << ", ";
  }
  ss << ")";
  // 打印操作数类型
  if (!operands.empty()) ss << " : (";
  // 然后打印operand
  for (auto [idx, opd] : IndexRange(operands)) {
    ss << opd->getValue()->getType()->toString();
    if (idx != operands.size() - 1) ss << ", ";
  }
  if (!operands.empty()) ss << ")";

  // 打印返回值类型
  if (!rets.empty()) ss << " -> ";
  for (auto [idx, ret] : IndexRange(rets)) {
    ss << ret->getType()->toString();
    if (idx != rets.size() - 1) ss << ", ";
  }
  return ss.str();
}
std::string ValuePrinter::print() { return print(root); }

}  // namespace lcir